# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
SETTLE Constraint algorithm
"""

import numpy as np
import mindspore as ms
import mindspore.numpy as msnp
from mindspore import ops, Tensor

from . import Constraint
from ...system import Molecule


class SETTLE(Constraint):
    """
    SETTLE constraint controller.

    `SETTLE: An Analytical Version of the SHAKE and RATTLE Algorithm for Rigid Water Models.`
        by Shuichi Miyamoto and Peter A. Kollman.

    Args:
        system (Molecule):          Simulation system.

    Inputs:
        - **coordinate** (Tensor) - The coordinates of the system.
        - **velocity** (Tensor) - The velocity of the system.
        - **force** (Tensor) - The force of the system.
        - **energy** (Tensor) - The energy of the system.
        - **kinetics** (Tensor) - The kinetics of the system.
        - **virial** (Tensor) - The virial of the system. Default: None
        - **pbc_box** (Tensor) - PBC box of the system. Default: None
        - **step** (int) - The step of the system. Default: 0

    Return:
        - coordinate (Tensor), Tensor of shape (B, A, D). Data type is float.
        - velocity (Tensor), Tensor of shape (B, A, D). Data type is float.
        - force (Tensor), Tensor of shape (B, A, D). Data type is float.
        - energy (Tensor), Tensor of shape (B, 1). Data type is float.
        - kinetics (Tensor), Tensor of shape (B, D). Data type is float.
        - virial (Tensor), Tensor of shape (B, D). Data type is float.
        - pbc_box (Tensor), Tensor of shape (B, D). Data type is float.

    Supported Platforms:
        ``Ascend`` ``GPU``
    """

    def __init__(self,
                 system: Molecule,
                 index: Tensor = None, # (B, X, 3)
                 ):
        super(SETTLE, self).__init__(system, bonds='all-bonds')
        # pylint: disable=invalid-name
        batches = self._coordinate.shape[0]
        atoms = self._coordinate.shape[1]

        # (B, G, D)
        self.target_yvector = Tensor(np.array([[[0, 1, 0]] * int(atoms/3)] * batches, np.float32))
        self.target_zvector = Tensor(np.array([[[0, 0, 1]] * int(atoms/3)] * batches, np.float32))

        # (B, 1, 4)
        self.identity_quaternion = Tensor([[[1, 0, 0, 0]]] * batches, ms.float32)
        self.inverse_quaternion = Tensor([[[1, -1, -1, -1]]] * batches, ms.float32)

        hx = 0.081649043
        hy = 0.057735897
        self.om = 15.9994
        self.hm = 1.0

        # (B, G, a, 1)
        self._mass = Tensor(np.array([[[self.om, self.hm, self.hm]] * int(atoms/3)] * batches, np.float32))[..., None]

        self.rb = hy / (1 + 2 * self.hm / self.om)
        self.ra = hy - self.rb
        self.rc = hx

        # (B, G, a, D)
        self.crd_1 = Tensor([[[0, self.ra, 0],
                              [-self.rc, -self.rb, 0],
                              [self.rc, -self.rb, 0]] * int(atoms/3)] * batches, ms.float32)
        self.aindices = Tensor([0, 1], ms.int32)
        self.bindices = Tensor([0, 2], ms.int32)
        self.cindices = Tensor([0, 3], ms.int32)
        self.hindices = Tensor([1, 2], ms.int32)

        self.dot = ops.Einsum('ijk,ijk->ij')


    def get_vector_transform(self, vec1, vec2):
        """ Get the transform quaternion from a vector to another.
        Args:
            vec1 (Tensor): The initial vector.
            vec2 (Tensor): The target vector.
        Returns:
            q (Tensor): The transform quaternion.
        """
        # (B, G, D)
        cross_vector = msnp.cross(vec1, vec2, axisc=-1)
        # (B, G, 1)
        cross_norm = msnp.norm(cross_vector, axis=-1, keepdims=True)
        zero_index = msnp.where(cross_norm == 0, 1, 0)
        qs = msnp.zeros((vec1.shape[0], vec1.shape[0], 1))
        # (B, G, D)
        qv = msnp.zeros_like(vec1)
        # (B, G, 1)
        qs += msnp.norm(vec1, axis=-1, keepdims=True) * msnp.norm(vec2, axis=-1, keepdims=True)
        qs += self.dot((vec1, vec2))[..., None]
        # (B, G, D)
        qv += cross_vector
        # (B, G, 4)
        q = msnp.concatenate((qs, qv), axis=-1)
        q = q * (1 - zero_index) + self.identity_quaternion * zero_index
        q /= msnp.norm(q, axis=-1, keepdims=True)
        return q

    def get_inverse(self, quater):
        """ Get the inverse operation of a given quaternion.
        Args:
            quater (Tensor): The given quaternion.
        Returns:
            quater^{-1}
        """
        factor = msnp.norm(quater, axis=-1, keepdims=True) ** 2
        return quater * self.inverse_quaternion / factor

    def quaternion_multiply(self, tensor_1, tensor_2):
        """ Calculate the quaternion multiplication.
        Args:
            tensor_1 (Tensor): The first quaternion, if the size of last dimension is 3, it will be padded to 4 auto.
            tensor_2 (Tensor): The second quaternion, if the size of last dimension is 3, it will be padded to 4 auto.
        Returns:
            q (Tensor): The quaternion product of tensor_1 and tensor_2.
        """
        # (B, G, 4)
        if tensor_1.shape[-1] == 3:
            tensor_1 = msnp.pad(tensor_1, ((0, 0), (0, 0), (1, 0)), mode='constant', constant_value=0)
        if tensor_2.shape[-1] == 3:
            tensor_2 = msnp.pad(tensor_2, ((0, 0), (0, 0), (1, 0)), mode='constant', constant_value=0)
        # (B, G, 1)
        s_1 = tensor_1[..., [0]]
        s_2 = tensor_2[..., [0]]
        # (B, G, 3)
        v_1 = tensor_1[..., 1:]
        v_2 = tensor_2[..., 1:]
        # (B, G, 1)
        s = s_1 * s_2
        # (B, G, 1)
        d = self.dot((v_1, v_2))[..., None]
        s -= d
        # (B, G, 3)
        v = msnp.zeros_like(v_1)
        v += s_1 * v_2
        v += v_1 * s_2
        v += msnp.cross(v_1, v_2, axisc=-1)
        q = msnp.concatenate((s, v), axis=-1)
        return q

    def hamiltonian_product(self, q, v):
        """ Perform hamiltonian product.
        Args:
            q (Tensor): The transform quaternion.
            v (Tensor): The vector to be transformed.
        Returns:
            qvq^{-1}
        """
        # (B, G, 4)
        iq = self.get_inverse(q)
        op1 = self.quaternion_multiply(v, iq)
        res = self.quaternion_multiply(q, op1)
        return res

    def group_hamiltonian_product(self, q, vec):
        """ Perform hamiltonian product in a 4-dimensional vector.
        Args:
            q (Tensor): The transform quaternion.
            vec (Tensor): The vector to be transformed.
        Returns:
            qvq^{-1}
        """
        # (B, G, 1, 4)
        group_q = q[..., None, :]
        # (B, G, a, 4)
        pad_q = msnp.pad(group_q, ((0, 0), (0, 0), (0, 2), (0, 0)), mode='wrap')
        # (B, G*a, 4)
        res_q = pad_q.reshape((q.shape[0], q.shape[1]*3, q.shape[2]))
        iq = self.get_inverse(res_q)
        op1 = self.quaternion_multiply(vec.reshape((q.shape[0], q.shape[1]*3, vec.shape[2])), iq)
        # (B, G, a, D)
        res = self.quaternion_multiply(res_q, op1)[..., 1:].reshape(vec.shape)
        return res


    def get_transform(self, crd_):
        """ Get the transform between A0B0C0 and a0b0c0.
        Args:
            crd_ (Tensor): The given water molecular coordinates.
        Returns:
            This function will return the transform and inverse transform from crd_ SETTLE axis.
        """
        # (B, G, A, D)
        crd = crd_.reshape((crd_.shape[0], -1, 3, crd_.shape[-1]))
        # (B, G, D)
        mass_center = msnp.sum(crd * self._mass, axis=-2) / self._mass.sum(axis=-2)
        # (B, G, D)
        cross_vector_1 = crd[..., 1, :] - crd[..., 0, :]
        cross_vector_2 = crd[..., 2, :] - crd[..., 1, :]
        cross_vector = msnp.cross(cross_vector_1, cross_vector_2, axisc=-1)
        cross_vector /= msnp.norm(cross_vector, axis=-1, keepdims=True)
        # (B, G, 4)
        transform_1 = self.get_vector_transform(cross_vector, self.target_zvector)

        inverse_transform_1 = self.get_inverse(transform_1)
        # (B, G, D)
        oa_vector = crd[..., 0, :] - mass_center
        oa_vector /= msnp.norm(oa_vector, axis=-1, keepdims=True)
        transformed_oa_vector = self.hamiltonian_product(transform_1, oa_vector)[..., 1:]
        transformed_oa_vector /= msnp.norm(transformed_oa_vector, axis=-1, keepdims=True)
        # (B, G, 4)
        transform_2 = self.get_vector_transform(transformed_oa_vector, self.target_yvector)
        inverse_transform_2 = self.get_inverse(transform_2)
        # (B, G, 4)
        combine_transform = self.quaternion_multiply(transform_2, transform_1)
        combine_itransform = self.quaternion_multiply(inverse_transform_1, inverse_transform_2)
        return combine_transform, combine_itransform

    def apply_transform(self, q, vec):
        """ Apply the quaternion transform q to a given vector vec. """
        # (B, G, a, D)
        crd = vec.reshape((vec.shape[0], -1, 3, vec.shape[-1]))
        crd -= self.get_mass_center(crd)
        # (B, G, a, 1)
        norm_crd = msnp.norm(crd, axis=-1, keepdims=True)
        # (B, G, a, D)
        crd /= norm_crd
        transform_crd = self.group_hamiltonian_product(q, crd)
        res = transform_crd * norm_crd
        # (B, A, D)
        return res.reshape(vec.shape)

    def _rotation(self, crd_1_0_, crd_1_1_):
        """ Get the rotation quaternion between triangles. """
        crd_1_0 = crd_1_0_.reshape((crd_1_0_.shape[0], -1, 3, crd_1_0_.shape[-1]))
        crd_1_1 = crd_1_1_.reshape((crd_1_1_.shape[0], -1, 3, crd_1_1_.shape[-1]))
        # (B, G)
        sin_phi = crd_1_1[..., 0, 2] / self.ra
        sin_phi = sin_phi.clip(-0.999999, 0.999999)
        phi = msnp.arcsin(sin_phi)
        sin_psi = (crd_1_1[..., 1, 2] - crd_1_1[..., 2, 2]) / 2 / self.rc / msnp.sqrt(1 - sin_phi ** 2)
        sin_psi = sin_psi.clip(-0.999999, 0.999999)
        psi = msnp.arcsin(sin_psi)
        xb2 = -self.rc * msnp.sqrt(1 - sin_psi ** 2)
        yb2 = -self.rb * msnp.sqrt(1 - sin_phi ** 2) - self.rc * sin_phi * sin_psi
        yc2 = -self.rb * msnp.sqrt(1 - sin_phi ** 2) + self.rc * sin_phi * sin_psi
        alpha = xb2 * (crd_1_0[..., 1, 0] - crd_1_0[..., 2, 0]) + \
                yb2 * (crd_1_0[..., 1, 1] - crd_1_0[..., 0, 1]) + \
                yc2 * (crd_1_0[..., 2, 1] - crd_1_0[..., 0, 1])
        beta = xb2 * (crd_1_0[..., 2, 1] - crd_1_0[..., 1, 1]) + \
               yb2 * (crd_1_0[..., 1, 0] - crd_1_0[..., 0, 0]) + \
               yc2 * (crd_1_0[..., 2, 0] - crd_1_0[..., 0, 0])
        gamma = (crd_1_0[..., 1, 0] - crd_1_0[..., 0, 0]) * crd_1_1[..., 1, 1] - \
                crd_1_1[..., 1, 0] * (crd_1_0[..., 1, 1] - crd_1_0[..., 0, 1]) + \
                (crd_1_0[..., 2, 0] - crd_1_0[..., 0, 0]) * crd_1_1[..., 2, 1] - \
                crd_1_1[..., 2, 0] * (crd_1_0[..., 2, 1] - crd_1_0[..., 0, 1])
        sin_theta = (alpha * gamma - beta * msnp.sqrt(alpha ** 2 + beta ** 2 - gamma ** 2)) / (alpha ** 2 + beta ** 2)
        sin_theta = sin_theta.clip(-0.999999, 0.999999)
        theta = msnp.arcsin(sin_theta)

        # (B, G, 4)
        cos_phi = msnp.pad(msnp.cos(phi[..., None] / 2), ((0, 0), (0, 0), (0, 3)), mode='constant', constant_value=0)
        sin_phi = msnp.pad(msnp.sin(phi[..., None] / 2), ((0, 0), (0, 0), (0, 3)), mode='constant', constant_value=0)
        cos_psi = msnp.pad(msnp.cos(psi[..., None] / 2), ((0, 0), (0, 0), (0, 3)), mode='constant', constant_value=0)
        sin_psi = msnp.pad(msnp.sin(psi[..., None] / 2), ((0, 0), (0, 0), (0, 3)), mode='constant', constant_value=0)
        cos_theta = msnp.pad(msnp.cos(theta[..., None] / 2), ((0, 0), (0, 0), (0, 3)), mode='constant',
                             constant_value=0)
        sin_theta = msnp.pad(msnp.sin(theta[..., None] / 2), ((0, 0), (0, 0), (0, 3)), mode='constant',
                             constant_value=0)

        phi_minus_value = -sin_phi[..., self.aindices]
        psi_minus_value = -sin_psi[..., self.bindices]
        theta_minus_value = -sin_theta[..., self.cindices]
        phi_add_value = sin_phi[..., self.aindices[::-1]]
        psi_add_value = sin_psi[..., self.bindices[::-1]]
        theta_add_value = sin_theta[..., self.cindices[::-1]]

        sin_phi[..., self.aindices] += phi_minus_value
        sin_phi[..., self.aindices] += phi_add_value
        sin_psi[..., self.bindices] += psi_minus_value
        sin_psi[..., self.bindices] += psi_add_value
        sin_theta[..., self.cindices] += theta_minus_value
        sin_theta[..., self.cindices] += theta_add_value

        # (B, 4)
        quater_phi = cos_phi + sin_phi
        quater_psi = cos_psi + sin_psi
        quater_theta = cos_theta + sin_theta

        rotation = self.quaternion_multiply(quater_theta, self.quaternion_multiply(quater_phi, quater_psi))
        rotation /= msnp.norm(rotation, axis=-1, keepdims=True)
        return rotation

    def _swap_h(self, crd_1_1_, crd_1_3_):
        """ Swap the H atoms if the resulted hydrogen from SETTLE are mixed. """
        # (B, G, a, D)
        #pylint: disable=invalid-name
        crd_1_1 = crd_1_1_.reshape((crd_1_1_.shape[0], -1, 3, crd_1_1_.shape[-1]))
        crd_1_3 = crd_1_3_.reshape((crd_1_3_.shape[0], -1, 3, crd_1_3_.shape[-1]))
        # (B, G)
        Hs_0_0 = crd_1_1[..., 1, 2]
        Hs_0_1 = crd_1_1[..., 2, 2]
        mask_0 = Hs_0_0 >= Hs_0_1
        Hs_1_0 = crd_1_3[..., 1, 2]
        Hs_1_1 = crd_1_3[..., 2, 2]
        mask_1 = Hs_1_0 >= Hs_1_1
        # (B, G, 1, 1)
        H_swap_mask = msnp.where(mask_0 == mask_1, 0, 1)[..., None, None]
        # (B, G, 2, D)
        h_minus_value = -crd_1_3[..., self.hindices, :] * H_swap_mask
        h_add_value = crd_1_3[..., self.hindices[::-1], :] * H_swap_mask
        # Swap the coordinate of hydrogen
        crd_1_3[..., self.hindices, :] += h_minus_value
        crd_1_3[..., self.hindices, :] += h_add_value

        new_crd = crd_1_3.reshape(crd_1_3_.shape)
        return new_crd

    def get_mass_center(self, crd_):
        """ Get the mass center of a given molecule. """
        # (B, G, A, D)
        crd = crd_.reshape((crd_.shape[0], -1, 3, crd_.shape[-1]))
        # (B, G, 1, D)
        cm = msnp.sum(crd * self._mass, axis=-2) / self._mass.sum(axis=-2)
        return cm[..., None, :]

    def get_vel_force_update(self, crd0_, vel0_):
        """ Get the update of velocity and force.
        Args:
            crd0_ (Tensor): The coordinate after SETTLE in the origin axis.
            vel0_ (Tensor): The initial velocity.
        Returns:
            dv (Tensor): The constraint velocity.
            df (Tensor): The constraint force.
        """
        #pylint: disable=invalid-name
        crd0 = crd0_.reshape((crd0_.shape[0], -1, 3, 3))
        vel0 = vel0_.reshape((vel0_.shape[0], -1, 3, 3))
        # (B, G, 1)
        AB = msnp.norm(crd0[..., 1, :] - crd0[..., 0, :], axis=-1, keepdims=True)
        AC = msnp.norm(crd0[..., 2, :] - crd0[..., 0, :], axis=-1, keepdims=True)
        BC = msnp.norm(crd0[..., 2, :] - crd0[..., 1, :], axis=-1, keepdims=True)

        # (B, G, D)
        e_AB = (crd0[..., 1, :] - crd0[..., 0, :]) / AB
        e_CA = (crd0[..., 0, :] - crd0[..., 2, :]) / AC
        e_BC = (crd0[..., 2, :] - crd0[..., 1, :]) / BC

        # (B, G)
        cosA = (AB ** 2 + AC ** 2 - BC ** 2) / 2 / AB / AC
        cosB = (AB ** 2 + BC ** 2 - AC ** 2) / 2 / AB / BC
        cosC = (BC ** 2 + AC ** 2 - AB ** 2) / 2 / BC / AC

        d = self.time_step * (2 * (self.om + self.hm) ** 2 + 2 * self.om * self.hm * cosA * cosB * cosC -
                              2 * (self.hm ** 2) * (cosA ** 2) - self.om * (self.om + self.hm) * (cosB ** 2 +
                                                                                                  cosC ** 2)) /\
            2 / self.hm
        # (B, G, D)
        v_AB_ = vel0[..., 1, :] - vel0[..., 0, :]
        v_CA_ = vel0[..., 0, :] - vel0[..., 2, :]
        v_BC_ = vel0[..., 2, :] - vel0[..., 1, :]

        # (B, G, 1)
        v_AB = self.dot((e_AB, v_AB_))[..., None]
        v_CA = self.dot((e_CA, v_CA_))[..., None]
        v_BC = self.dot((e_BC, v_BC_))[..., None]

        tau_AB = self.om * (v_AB * (2 * (self.om + self.hm) - self.om * cosC ** 2) +
                            v_BC * (self.hm * cosC * cosA - (self.om + self.hm) * cosB) +
                            v_CA * (self.om * cosB * cosC - 2 * self.hm * cosA)) / d
        tau_BC = (v_BC * ((self.om + self.hm) ** 2 - (self.hm ** 2) * (cosA ** 2)) +
                  v_CA * self.om * (self.hm * cosA * cosB - (self.om + self.hm) * cosC) +
                  v_AB * self.om * (self.hm * cosC * cosA - (self.om + self.hm) * cosB)) / d
        tau_CA = self.om * (v_CA * (2 * (self.om + self.hm) - self.om * cosB ** 2) +
                            v_AB * (self.om * cosB * cosC - 2 * self.hm * cosA) +
                            v_BC * (self.hm * cosA * cosB - (self.om + self.hm) * cosC)) / d

        # (B, G, D)
        df_A = tau_AB * e_AB - tau_CA * e_CA
        dv_A = self.time_step * df_A / 2 / self.om
        df_B = tau_BC * e_BC - tau_AB * e_AB
        dv_B = self.time_step * df_B / 2 / self.hm
        df_C = tau_CA * e_CA - tau_BC * e_BC
        dv_C = self.time_step * df_C / 2 / self.hm

        # (B, A, D)
        return msnp.stack((dv_A, dv_B, dv_C), axis=-2).reshape(vel0_.shape), \
            msnp.stack((df_A, df_B, df_C), axis=-2).reshape(vel0_.shape)

    def construct(self,
                  coordinate: Tensor,
                  velocity: Tensor,
                  force: Tensor,
                  energy: Tensor,
                  kinetics: Tensor,
                  virial: Tensor = None,
                  pbc_box: Tensor = None,
                  step: int = 0,
                  ):
        #pylint: disable=invalid-name
        # (B, A, D)
        crd_old_0_0 = self._coordinate.copy()
        crd_new_0_1 = coordinate.copy()
        mc = self.get_mass_center(coordinate)
        # (B, G, 4)
        transform, itransform = self.get_transform(crd_old_0_0)
        # (B, A, D)
        crd_old_1_0 = self.apply_transform(transform, crd_old_0_0)
        crd_new_1_1 = self.apply_transform(transform, crd_new_0_1)
        rotation = self._rotation(crd_old_1_0, crd_new_1_1)
        crd_new_1_3 = self.apply_transform(rotation, self.crd_1)
        crd_new_1_3 = self._swap_h(crd_new_1_1, crd_new_1_3)
        # (B, G, a, D)
        crd_new_0_3 = self.apply_transform(itransform, crd_new_1_3).reshape((crd_old_0_0.shape[0], -1, 3,
                                                                             crd_old_0_0.shape[-1]))
        # (B, G, a, D)
        crd_new_0_3 += mc
        # (B, A, D)
        crd_new_0_3 = crd_new_0_3.reshape(crd_old_0_0.shape)

        dv, df = self.get_vel_force_update(coordinate, velocity)
        coordinate = crd_new_0_3

        velocity += dv
        force += df

        return coordinate, velocity, force, energy, kinetics, virial, pbc_box
